#include <sys/time.h> 
#include <iostream>

#include "randomc.h"
#include "stocc.h"

#ifndef MULTIFILE_PROJECT
// If compiled as a single file then include these cpp files, 
// If compiled as a project then compile and link in these cpp files.
   #include "mersenne.cpp"             // code for random number generator
   #include "stoc1.cpp"                // random library source code
   #include "userintf.cpp"             // define system specific user interface
#endif

using namespace std;

// functions
void Index0(const int *Y, int *N0, const int N)
{
  int count = 0;
  for (int n = 0; n < N; n++){
    if (Y[n] == 0){
      N0[count] = n;
      count++;
    }
  }
}

void SumSumSq(const float *X, float *Sum, float *SumSq, const int N, const int K){
  for (int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;
    for (int n = 0; n < N; n++){
      Sum[k] += X[iter+n];
      SumSq[k] += X[iter+n] * X[iter+n];
    }
  }
}

void Yperm(int *Y, const int N)
{
  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  for (int n = 1; n < N; n++){
    int idx = sto.IRandom(0, n); // uniform over [0,n]
    float val = Y[n];
    Y[n] = Y[idx];
    Y[idx] = val;
  }
}

void Tstat(const float *X, float *T, const float *Sum, const float *SumSq, const int *N0, const int N, const int K, const int n0, const int n1)
{
  float s0;
  float s1;

  float x0;
  float x1;

  float xj;
  
  for(int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    s0 = 0.0;
    x0 = 0.0;
    for (int n = 0; n < n0; n++){
      xj = X[iter + N0[n]];
      x0 += xj;
      s0 += xj * xj;
    }
    x0 /= n0;
    x1  = (Sum[k] - n0*x0) / n1;
     
    s1  = (SumSq[k]-s0-n1*x1*x1) / (n1-1);
    s0 = (s0-n0*x0*x0) / (n0-1);

    T[k] = (x0-x1) / sqrt(s0/n0+s1/n1);
  }
}

void TstatTmax(const float *X, float *Ttilde, const float *Sum, const float *SumSq, const int *N0, const int N, const int K, const int n0, const int n1, float *Tmax)
{
  float s0;
  float s1;

  float x0;
  float x1;

  float xj;

  float zeta = 0.0;
  
  for(int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    s0 = 0.0;
    x0 = 0.0;
    for (int n = 0; n < n0; n++){
      xj = X[iter + N0[n]];
      x0 += xj;
      s0 += xj * xj;
    }
    x0 /= n0;
    x1  = (Sum[k] - n0*x0) / n1;
     
    s1  = (SumSq[k]-s0-n1*x1*x1) / (n1-1);
    s0 = (s0-n0*x0*x0) / (n0-1);

    Ttilde[k] = (x0-x1) / sqrt(s0/n0+s1/n1);

    if (zeta < fabs(Ttilde[k]))
      zeta = fabs(Ttilde[k]);
  }
  Tmax[0] = zeta;
}

void Pvalue(const float *T, const float *Ttilde, float *p, float *P, const float *Tmax, const int K)
{
  for(int k = 0; k < K; k++){
    if(fabs(Ttilde[k]) >= fabs(T[k]))
      p[k]++;
    if(Tmax[0] >= fabs(T[k]))
      P[k]++;
  }
}

long int seed = (long int)time(0); // global seed

int main(const int argc, const char *argv[])
{
  // argv[1]: N
  // argv[2]: K
  // argv[3]: B

  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  int N = atoi(argv[1]);
  int K = atoi(argv[2]);
  int B = atoi(argv[3]);

  // print
  std::cout<<"(N="<<N<<",K="<<K<<",B="<<B<<")"<<"\n";

  // allocate memory  
  float *X;
  X = (float *)malloc(N*K*sizeof(float));
  if (X == NULL) exit (1);

  int *Y;
  Y = (int *)malloc(N*sizeof(int));
  if (Y == NULL) exit (1);

  float *p;
  p = (float *)malloc(K*sizeof(float));
  if (p == NULL) exit (1);

  float *P;
  P = (float *)malloc(K*sizeof(float));
  if (P == NULL) exit (1);

  float *T;
  T = (float *)malloc(K*sizeof(float));
  if (T == NULL) exit (1);

  float *Ttilde;
  Ttilde = (float *)malloc(K*sizeof(float));
  if (Ttilde == NULL) exit (1);

  float *Sum;
  Sum = (float *)malloc(K*sizeof(float));
  if (Sum == NULL) exit (1);

  float *SumSq;
  SumSq = (float*)malloc(K*sizeof(float));
  if (SumSq == NULL) exit (1);

  float *Tmax;
  Tmax = (float *)malloc(sizeof(float));
  if (Tmax == NULL) exit (1);
  Tmax[0] = 0.0;

  int count = 0;
  double prob_event = 0.5; // probability of event
  for (int n = 0; n < N; n++){
    Y[n] = sto.Bernoulli(prob_event);
    for (int k = 0; k < K; k++){
      X[count] = sto.Normal(0, 1.0);
      count++;
    }
  }

  // start time profiling
  timeval t1, t2;
  double elapsed;
  gettimeofday(&t1, NULL);

  for (int k = 0; k < K; k++){
    p[k] = 0.0;
    P[k] = 0.0;
    T[k] = 0.0;
    Ttilde[k] = 0.0;
    Sum[k] = 0.0;
    SumSq[k] = 0.0;
  }

  // compute n0 and n1
  int n0 = 0;
  for (int n = 0; n < N; n++){
    if (Y[n] == 0)
      n0++;
  }
  int n1 = N - n0;

  // compute N0
  int *N0;
  N0 = (int *)malloc(n0*sizeof(int));
  Index0(Y, N0, N);

  SumSumSq(X, Sum, SumSq, N, K);

  // compute T
  Tstat(X, T, Sum, SumSq, N0, N, K, n0, n1);

  for (int b = 0; b < B; b++){
    // permute Y
    Yperm(Y, N);

    // compute N0
    Index0(Y, N0, N);

    // compute test statistic and find the maximum
    TstatTmax(X, Ttilde, Sum, SumSq, N0, N, K, n0, n1, Tmax);

    // compute p-values
    Pvalue(T, Ttilde, p, P, Tmax, K);
  }

  gettimeofday(&t2, NULL);
  elapsed = (t2.tv_sec - t1.tv_sec) * 1000.0;      // sec to ms
  elapsed += (t2.tv_usec - t1.tv_usec) / 1000.0;

  std::cout<<"elapsed time:"<<elapsed<<" milliseconds"<<"\n";

  // cleanup
  free(Y);
  free(X);
  free(p);
  free(P);
  free(T);
  free(Ttilde);
  free(SumSq);
  free(Sum);
  free(Tmax);
  free(N0);

  return(0);
}
